import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split


class CompositionPredictor:
    def __init__(self, excel_path, performance_columns, composition_columns, max_iterations=20, tolerance=0.1):
        self.excel_path = excel_path
        self.performance_columns = performance_columns
        self.composition_columns = composition_columns
        self.max_iterations = max_iterations
        self.tolerance = tolerance
        self.data = self.load_data()

    def load_data(self):
        # 读取 Excel 文件
        try:
            data = pd.read_excel(self.excel_path)
            print("数据读取成功")
            return data
        except Exception as e:
            print(f"读取数据时发生错误: {e}")
            exit()

    def train_random_forest(self, X, y):
        # 划分数据集
        x_train, x_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=None)

        model = RandomForestRegressor(n_estimators=300, oob_score=True, min_samples_leaf=1)
        model.fit(x_train, y_train)

        # 计算模型在测试集上的得分
        score = model.score(x_test, y_test)
        print(f"模型在测试集上的得分: {score:.4f}")

        return model

    def predict_composition(self, input_performance):
        within_tolerance = False

        # 迭代训练和预测
        for iteration in range(self.max_iterations):
            # P-C 模型训练：性能预测成分
            X_pc = self.data[self.performance_columns]
            y_pc = self.data[self.composition_columns]

            pc_model = self.train_random_forest(X_pc, y_pc)

            # 使用 P-C 模型预测成分
            predicted_composition = pc_model.predict(input_performance)
            np.set_printoptions(suppress=True, precision=4)
            predicted_composition_df = pd.DataFrame(predicted_composition, columns=self.composition_columns)

            # 使用 C-P 模型训练：成分预测性能
            predicted_performance = []
            for performance in self.performance_columns:
                y_cp = self.data[[performance]].values.ravel()
                cp_model = self.train_random_forest(self.data[self.composition_columns], y_cp)
                predicted_value = cp_model.predict(predicted_composition_df)
                predicted_performance.append(predicted_value[0])

            # 转换为数组便于比较
            predicted_performance = np.array(predicted_performance).reshape(1, -1)

            # 计算误差并检查是否在容忍范围内
            differences = abs((input_performance.values - predicted_performance) / input_performance.values)
            within_tolerance = np.all(differences < self.tolerance)

            # 打印每次迭代的预测信息
            print(f"迭代 {iteration + 1}:")
            print("预测成分:", predicted_composition)
            print("预测性能:", predicted_performance)
            print("误差百分比:", differences)
            print("")

            # 如果在容忍范围内，输出结果并退出循环
            if within_tolerance:
                print(f"在第 {iteration + 1} 次迭代时找到了符合的成分组合:")
                print("预测成分:", predicted_composition)
                print("预测性能:", predicted_performance)
                break

        # 如果没有找到符合条件的组合
        if not within_tolerance:
            print("未找到符合误差范围的成分组合，请调整容忍范围或增加迭代次数。")

def main():
    # 实例化和使用类
    excel_path = r'C:\Users\13945\Desktop\MLDS铝合金成分设计数据.xlsx'
    performance_columns = ["Ultimate tensile strength/MPa", "Fracture thougness/MPa.m1/2", "Elongation/%"]
    composition_columns = ['Si/％', 'Mn/％', 'Zn/％', 'Mg/％', 'Cu/％', 'Cr/％', 'Zr/％', 'Ti/％', 'Fe/％', 'Ni/％', 'other']
    max_iterations = 20
    tolerance = 0.1
    # 输入性能列表（示例）
    input_performance = pd.DataFrame([[590, 44, 16.5]], columns=performance_columns)

    # 创建实例并进行预测   excel_path, performance_columns, composition_columns, max_iterations=20, tolerance=0.1
    predictor = CompositionPredictor(excel_path=excel_path,
                                     performance_columns=performance_columns,
                                     composition_columns=composition_columns,
                                     max_iterations=max_iterations,
                                     tolerance=tolerance)
    predictor.predict_composition(input_performance)


if __name__ == "__main__":
    main()